package org.example.offical.doc.rag;

import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.model.openai.OpenAiChatModelName;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.aggregator.DefaultContentAggregator;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.content.retriever.WebSearchContentRetriever;
import dev.langchain4j.rag.query.router.DefaultQueryRouter;
import dev.langchain4j.rag.query.transformer.ExpandingQueryTransformer;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.Result;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.service.V;
import dev.langchain4j.service.tool.ToolExecution;
import dev.langchain4j.store.embedding.filter.MetadataFilterBuilder;
import dev.langchain4j.store.memory.chat.InMemoryChatMemoryStore;
import dev.langchain4j.web.search.searchapi.SearchApiWebSearchEngine;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.example.offical.doc.ChatMemoryDemo;
import org.example.offical.doc.ModelUtils;
import org.example.offical.doc.ai.service.AiServiceWithMemoryAndToolExecution;
import org.junit.Test;

import java.util.List;

/**
 * @author superMan
 * @since fish_temp_since
 */
@Slf4j
public class AdvanceRAG {
    private static final String SEARCHAPI_API_KEY = "eCdrsTuv3tpQYJ7vaSGyNUsJ";

    public static void main(String[] args) {


        /*
        userMessage
        Query
        多个query QueryRouter
        对应多个ContentRetrievers 内容检索
        将所有检索结果合并最终排名列表 ContentAggregator -> Content
        将上面生成的Content注入到用户的UserMessage上下文
        最终，包含用户的UserMessage和检索到的内容，一起发送给LLM，LLM根据上下文进行回答。
         */
        // RetrievalAugmentor是RAG流水线的入口点。它负责用从各种来源检索到的相关ChatMessage来扩充Content

        EmbeddingStoreContentRetriever embeddingStoreContentRetriever = embeddingStoreContentRetriever();
        WebSearchContentRetriever webSearchContentRetriever = webSearchContentRetriever();
        PromptTemplate promptTemplate = promptTemplate();

        // 通过QueryRouter可以实现多个检索（contentRetriever）的组合
        DefaultQueryRouter queryRouter = new DefaultQueryRouter(embeddingStoreContentRetriever, webSearchContentRetriever);

        Assistant assistant = AiServices.builder(Assistant.class)
                // 检索增强
                .retrievalAugmentor(DefaultRetrievalAugmentor.builder()
                        // 查询转换器，能够对用户输入进行优化
                        //.queryTransformer(new ExpandingQueryTransformer(ModelUtils.getOpenAiDemoModel()))
                        // 检索
                        // .contentRetriever(embeddingStoreContentRetriever)
                        // 通过QueryRouter实现多个检索（contentRetriever）的组合
                        .queryRouter(queryRouter)
                        // 内容聚合，重排序
                        .contentAggregator(new DefaultContentAggregator())
                        // 内容注入
                        .contentInjector(DefaultContentInjector.builder().promptTemplate(promptTemplate).build())
                        // 如果通过queryRouter使用了多个召回，则可以使用多线程
                        //.executor(Runnable::run)
                        .build())
                // 记忆
                .chatMemoryProvider(memoryId -> TokenWindowChatMemory.builder()
                        .maxTokens(1000, new OpenAiTokenizer(OpenAiChatModelName.GPT_4_O_MINI))
                        .chatMemoryStore(new InMemoryChatMemoryStore())
                        .build())
                // 使用的大语言模型
                .chatLanguageModel(ModelUtils.getOpenAiDemoModel())
                .tools(new AiServiceWithMemoryAndToolExecution.Tools())
                .build();

        // Result<String> chatRet = assistant.chat("请自我介绍一下");
        Result<String> chatRet = assistant.chat("三国演义中的人物");
        // Result<String> chatRet = assistant.chat("张三的bmi");
        List<Content> sources = chatRet.sources();
        for (Content source : sources) {
            System.out.println("源引用：" + source);
        }
        List<ToolExecution> toolExecutions = chatRet.toolExecutions();
        if (CollectionUtils.isNotEmpty(toolExecutions)) {
            for (ToolExecution toolExecution : toolExecutions) {
                System.out.println("tool execution");
                System.out.println(toolExecution);
            }
        }
    }


    @Test
    public void addSth2RedisEmbeddingVectorStore() {
        Document document = FileSystemDocumentLoader.loadDocument("F:\\AI\\代码\\mine\\first-ai-loulan\\src\\main\\java\\org\\example\\offical\\doc\\rag\\users.txt", new ApacheTikaDocumentParser(true));
        List<TextSegment> textSegments = DocumentSplitters.recursive(30, 0)
                .split(document);
        Response<List<Embedding>> embedded = ModelUtils.getOpenAiEmbeddingDemoModel().embedAll(textSegments);
        RedisEmbeddingStore redisEmbeddingStore = redisEmbeddingStore();
        for (int i = 0; i < embedded.content().size(); i++) {
            Embedding embedding = embedded.content().get(i);
            TextSegment textSegment = textSegments.get(i);
            String vectorId = redisEmbeddingStore.add(embedding, textSegment);
            log.info("嵌入向量和文本片段到向量库: id: {}, 文本: {}, 向量:{}", vectorId, textSegment, embedded);
        }
    }

    /**
     * 提示词模板
     *
     * @return
     */
    private static PromptTemplate promptTemplate() {
        // PromptTemplate必须包含{{userMessage}}和{{contents}}变量
        return PromptTemplate.from("""
                用户的问题：{{userMessage}}
                
                **请严格根据上下文中的内容来回答用户的问题”。**
                请不要随便杜撰，如果上下文中的内容为空，请回答"我还在学习当中呢..."
                以下内容为上下文：
                
                {{contents}}
                """);
    }

    /**
     * 基于联网搜索的文本召回
     *
     * @return
     */
    private static WebSearchContentRetriever webSearchContentRetriever() {
        return WebSearchContentRetriever.builder()
                .webSearchEngine(SearchApiWebSearchEngine.builder()
                        .apiKey(SEARCHAPI_API_KEY)
                        .build())
                .maxResults(3)
                .build();
    }


    /**
     * 基于向量库的文本召回
     *
     * @return
     */
    private static EmbeddingStoreContentRetriever embeddingStoreContentRetriever() {
        RedisEmbeddingStore redisEmbeddingStore = redisEmbeddingStore();
        return EmbeddingStoreContentRetriever.builder()
                // 向量库的实现
                .embeddingStore(redisEmbeddingStore)
                // 向量模型
                .embeddingModel(ModelUtils.getOpenAiEmbeddingDemoModel())
                .maxResults(3)
                .minScore(0.6)
                // 根据文档的元数据过滤数据
                // .filter(MetadataFilterBuilder.metadataKey("userId").isEqualTo("12345"))
                // filter can also be specified dynamically depending on the query
                // .dynamicFilter(query -> {
                //     // userService.queryUser();
                //     return MetadataFilterBuilder.metadataKey("userId").isEqualTo("123");
                // })
                .build();
    }

    /**
     * 基于redis的向量库
     *
     * @return
     */
    private static RedisEmbeddingStore redisEmbeddingStore() {
        return RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .prefix("advance-rag:")
                .indexName("advanceRag")
                .dimension(1536)
                .build();
    }


    interface Assistant {
        Result<String> chat(String userMessage);
    }
}
